import math
import random
import torch

class PoissonEncoding:
    HUGE_VAL = 0x77777777
    
    def __init__(self, rate = 100) -> None:
        self.rate = rate
        self.tn = 0
        self.T = 1000 / self.rate

    def setT(self, rate):
        self.rate = rate
        self.T = 1000 / rate
    
    def spike(self):
        u = int(self.update() * 10)
        while u < 10:
            u = int(self.update() * 10)
        self.tn += u * 0.1
        return self.tn

    def update(self):
        if self.T <= 0:
            uptime = self.HUGE_VAL
        else:
            uptime = -self.T * math.log(1.0 - random.random())
        return uptime
    
    def spike_encoding(self, duration) -> torch.Tensor:
        seq = []
        t = self.spike()
        while t <= duration:
            seq.append(round(t, 1))
            t = self.spike()
        return torch.Tensor(seq)
    

if __name__ == '__main__':
    pe = PoissonEncoding(50)
    s = pe.spike_encoding(100)
    print(s)